package com.hk.core.sharding.hibernate.jpa;

import com.hk.commons.multischema.DataSourceType;
import com.hk.commons.util.StringUtils;
import com.hk.core.data.jpa.repository.advice.RepositoryInterfaceAround;
import com.hk.core.sharding.HintToken;
import com.hk.core.sharding.ThreadLocalHintTokenProvider;
import org.aopalliance.intercept.MethodInvocation;

import java.lang.reflect.Method;
import java.util.Objects;

/**
 * @author Kevin
 * @date 2024-07-03 17:55
 */
public class ShardingDataSourceRepositoryAround implements RepositoryInterfaceAround {

    @Override
    public void before(MethodInvocation invocation, Class<?> repositoryClass) {
        String database = lookupForMethod(invocation.getMethod());
        if (StringUtils.hasLength(database)) {
            ThreadLocalHintTokenProvider.set(new HintToken().setDatabase(database));
        }
    }

    protected String lookupDatasource(String datasource) {
        return datasource;
    }

    private String lookupForMethod(Method method) {
        DataSourceType annotation = method.getAnnotation(DataSourceType.class);
        return Objects.nonNull(annotation) ? lookupDatasource(annotation.value()) : null;
    }

    @Override
    public void after(Class<?> repositoryClass) {
        ThreadLocalHintTokenProvider.clearWithTransaction();
    }
}
